package com.example.common.grpc;

import com.example.common.constants.ServiceNames;
import com.example.game.proto.internal.GrpcServiceGrpc;
import io.grpc.ConnectivityState;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * 基于服务器ID的负载均衡器
 * 支持多实例选择和健康检查
 *
 * @author: Allen
 * @create: 2025/7/3 22:35
 * @description: 服务器ID负载均衡器
 **/
@Slf4j
@Component
public class ServerIdLoadBalancer {

    // 存储每个服务类型的可用实例
    public final Map<ServiceNames, List<ServerInstance>> serviceInstances = new ConcurrentHashMap<>();

    // 轮询计数器
    private final Map<ServiceNames, AtomicInteger> roundRobinCounters = new ConcurrentHashMap<>();

    // 服务器ID到实例的映射
    private final Map<Integer, ServerInstance> serverIdToInstance = new ConcurrentHashMap<>();

    /**
     * 服务器实例信息
     */
    public static class ServerInstance {
        private final int serverId;
        private final ServiceNames serviceType;
        private volatile boolean healthy = true;
        //private List<RpcClient> channels;
        private GrpcServiceGrpc.GrpcServiceBlockingStub stub;
        private GrpcServiceGrpc.GrpcServiceStub asyncStub;
        private volatile long lastHealthCheck = System.currentTimeMillis();
        private final int weight;
        private int connNum = 5;
        private Random random;
        private ManagedChannel channel;

        public ServerInstance(int serverId, ServiceNames serviceType, SocketAddress socketAddress, int grpcKeepAliveTime,
                              int weight) {
            this.serverId = serverId;
            this.serviceType = serviceType;
         /*   this.channels = new ArrayList<>(connNum);
            for (int i = 0; i < connNum; i++) {
                RpcClient client = new RpcClient(serverId, serviceType, socketAddress, grpcKeepAliveTime);
                client.init();
                channels.add(client);
            }*/
            channel = createConservativeChannel((InetSocketAddress)socketAddress);

            this.stub =   GrpcServiceGrpc.newBlockingStub(channel);;
            this.asyncStub =  GrpcServiceGrpc.newStub(channel);;

            this.weight = weight;
            random = new Random();
        }

        public ServerInstance(int serverId, ServiceNames serviceType,
                              GrpcServiceGrpc.GrpcServiceBlockingStub stub,
                              GrpcServiceGrpc.GrpcServiceStub asyncStub,
                              ManagedChannel channel, int weight) {
            this.serverId = serverId;
            this.serviceType = serviceType;
            this.stub = stub;
            this.asyncStub = asyncStub;
            this.channel = channel;
            this.weight = weight;
        }

        public ManagedChannel createConservativeChannel(InetSocketAddress socketAddress) {
            return ManagedChannelBuilder.forAddress(socketAddress.getHostName(), socketAddress.getPort())
                    .usePlaintext()
                    // 保守的keepalive配置
                    .keepAliveTime(15, TimeUnit.SECONDS)     // 120秒发送一次keepalive ping
                    .keepAliveTimeout(15, TimeUnit.SECONDS)   // keepalive ping超时时间
                    // 消息大小配置
                    .maxInboundMessageSize(5 * 1024 * 1024)        // 512KB最大消息大小
                    // 连接配置
                    .build();
        }

        /*public RpcClient getClient() {
            int index = random.nextInt(connNum);
            RpcClient client = channels.get(index);
            return client;
        }*/

        public int getServerId() {
            return serverId;
        }

        public ServiceNames getServiceType() {
            return serviceType;
        }

        public GrpcServiceGrpc.GrpcServiceBlockingStub getStub() {
            return stub;
        }

        public GrpcServiceGrpc.GrpcServiceStub getAsyncStub() {
            return asyncStub;
        }

        public boolean isHealthy() {
            return healthy;
        }

        public int getWeight() {
            return weight;
        }

        public void setHealthy(boolean healthy) {
            this.healthy = healthy;
            this.lastHealthCheck = System.currentTimeMillis();
        }

        public long getLastHealthCheck() {
            return lastHealthCheck;
        }

        public ManagedChannel getChannel() {
            return channel;
        }

        @Override
        public String toString() {
            return String.format("ServerInstance{serverId=%d, serviceType=%s, healthy=%s, weight=%d}",
                    serverId, serviceType, healthy, weight);
        }
    }

    /**
     * 注册服务器实例
     */
    public void registerInstance(int serverId, ServiceNames serviceType,
                               GrpcServiceGrpc.GrpcServiceBlockingStub stub,
                               GrpcServiceGrpc.GrpcServiceStub asyncStub,
                               ManagedChannel channel) {
        registerInstance(serverId, serviceType, stub, asyncStub, channel, 1);
    }

    /**
     * 注册服务器实例（带权重）
     */
    public void registerInstance(int serverId, ServiceNames serviceType,
                               GrpcServiceGrpc.GrpcServiceBlockingStub stub,
                               GrpcServiceGrpc.GrpcServiceStub asyncStub,
                               ManagedChannel channel, int weight) {
        ServerInstance instance = new ServerInstance(serverId, serviceType, stub, asyncStub, channel, weight);

        serviceInstances.computeIfAbsent(serviceType, k -> new ArrayList<>()).add(instance);
        serverIdToInstance.put(serverId, instance);
        roundRobinCounters.computeIfAbsent(serviceType, k -> new AtomicInteger(0));

        log.info("注册服务器实例: {}", instance);
    }

    /**
     * 注销服务器实例
     */
    public void unregisterInstance(int serverId) {
        ServerInstance instance = serverIdToInstance.remove(serverId);
        if (instance != null) {
            serviceInstances.get(instance.getServiceType()).remove(instance);
            log.info("注销服务器实例: {}", instance);
        }
    }

    /**
     * 移除指定服务类型的所有实例
     */
    public void removeServiceInstances(ServiceNames serviceType) {
        List<ServerInstance> instances = serviceInstances.get(serviceType);
        if (instances != null) {
            // 从serverIdToInstance映射中移除
            for (ServerInstance instance : instances) {
                serverIdToInstance.remove(instance.getServerId());
            }
            // 清空该服务类型的实例列表
            instances.clear();
            log.info("移除所有{}服务实例，共{}个", serviceType, instances.size());
        }
    }

    /**
     * 根据服务器ID获取实例
     */
    public ServerInstance getInstanceByServerId(int serverId) {
        return serverIdToInstance.get(serverId);
    }

    /**
     * 根据服务类型获取实例（轮询负载均衡）
     */
    public ServerInstance getInstanceByServiceType(ServiceNames serviceType) {
        return getInstanceByServiceType(serviceType, LoadBalanceStrategy.ROUND_ROBIN);
    }

    /**
     * 根据服务类型和策略获取实例
     */
    public ServerInstance getInstanceByServiceType(ServiceNames serviceType, LoadBalanceStrategy strategy) {
        List<ServerInstance> instances = serviceInstances.get(serviceType);
        if (instances == null || instances.isEmpty()) {
            log.warn("没有可用的{}服务实例", serviceType);
            return null;
        }

        // 过滤健康的实例
        List<ServerInstance> healthyInstances = instances.stream()
                .filter(ServerInstance::isHealthy)
                .collect(Collectors.toList());

        if (healthyInstances.isEmpty()) {
            log.warn("没有健康的{}服务实例", serviceType);
            return null;
        }

        switch (strategy) {
            case ROUND_ROBIN:
                return getInstanceByRoundRobin(serviceType, healthyInstances);
            case WEIGHTED_ROUND_ROBIN:
                return getInstanceByWeightedRoundRobin(serviceType, healthyInstances);
            case RANDOM:
                return getInstanceByRandom(healthyInstances);
            case HASH:
                return getInstanceByHash(healthyInstances);
            default:
                return getInstanceByRoundRobin(serviceType, healthyInstances);
        }
    }

    /**
     * 轮询策略
     */
    private ServerInstance getInstanceByRoundRobin(ServiceNames serviceType, List<ServerInstance> instances) {
        AtomicInteger counter = roundRobinCounters.get(serviceType);
        int index = counter.getAndIncrement() % instances.size();
        return instances.get(index);
    }

    /**
     * 加权轮询策略
     */
    private ServerInstance getInstanceByWeightedRoundRobin(ServiceNames serviceType, List<ServerInstance> instances) {
        // 简化实现：按权重比例选择
        int totalWeight = instances.stream().mapToInt(ServerInstance::getWeight).sum();
        int randomWeight = new Random().nextInt(totalWeight);

        int currentWeight = 0;
        for (ServerInstance instance : instances) {
            currentWeight += instance.getWeight();
            if (randomWeight < currentWeight) {
                return instance;
            }
        }

        return instances.get(0); // 兜底
    }

    /**
     * 随机策略
     */
    private ServerInstance getInstanceByRandom(List<ServerInstance> instances) {
        int idx = new Random().nextInt(instances.size());
        return instances.get(idx);
    }

    /**
     * 哈希策略
     */
    private ServerInstance getInstanceByHash(List<ServerInstance> instances) {
        int hash = Objects.hash(System.currentTimeMillis(), Thread.currentThread().getId());
        return instances.get(Math.abs(hash) % instances.size());
    }

    /**
     * 健康检查
     */
    public void healthCheck() {
        for (List<ServerInstance> instances : serviceInstances.values()) {
            for (ServerInstance instance : instances) {
                try {
                    ConnectivityState state = instance.getChannel().getState(true);
                    boolean healthy = state == ConnectivityState.READY;
                    instance.setHealthy(healthy);

                    if (!healthy) {
                        log.warn("服务器实例不健康: {}", instance);
                    }
                } catch (Exception e) {
                    log.error("健康检查失败: {}", instance, e);
                    instance.setHealthy(false);
                }
            }
        }
    }

    /**
     * 获取服务实例统计信息
     */
    public Map<ServiceNames, ServiceStats> getServiceStats() {
        Map<ServiceNames, ServiceStats> stats = new HashMap<>();

        for (Map.Entry<ServiceNames, List<ServerInstance>> entry : serviceInstances.entrySet()) {
            ServiceNames serviceType = entry.getKey();
            List<ServerInstance> instances = entry.getValue();

            long totalInstances = instances.size();
            long healthyInstances = instances.stream().filter(ServerInstance::isHealthy).count();

            stats.put(serviceType, new ServiceStats(totalInstances, healthyInstances));
        }

        return stats;
    }

    /**
     * 服务统计信息
     */
    public static class ServiceStats {
        private final long totalInstances;
        private final long healthyInstances;

        public ServiceStats(long totalInstances, long healthyInstances) {
            this.totalInstances = totalInstances;
            this.healthyInstances = healthyInstances;
        }

        public long getTotalInstances() {
            return totalInstances;
        }

        public long getHealthyInstances() {
            return healthyInstances;
        }

        public double getHealthRate() {
            return totalInstances > 0 ? (double) healthyInstances / totalInstances : 0.0;
        }

        @Override
        public String toString() {
            return String.format("ServiceStats{total=%d, healthy=%d, healthRate=%.2f%%}",
                    totalInstances, healthyInstances, getHealthRate() * 100);
        }
    }

    /**
     * 负载均衡策略
     */
    public enum LoadBalanceStrategy {
        ROUND_ROBIN,      // 轮询
        WEIGHTED_ROUND_ROBIN, // 加权轮询
        RANDOM,           // 随机
        HASH              // 哈希
    }
} 